package com.seanLab.tool.TagSuggestModel;

import com.seanLab.domain.TagInfo;
import com.seanLab.dto.RecommendArticleDto;
import com.seanLab.dto.SuggestArticleKeywordsDto;
import com.seanLab.dto.SuggestModelArticleDto;
import com.seanLab.dto.SuggestModelImageDto;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

public class TagSuggestModelByER extends TagSuggestModel {

//    protected RawInfoType[] rawInfoNeed = new RawInfoType[]
//            {RawInfoType.TITLE, RawInfoType.CONTENT};
    protected RawInfoType[] rawInfoNeed = new RawInfoType[]
            {RawInfoType.CONTENT};
//    protected RawInfoType[] rawInfoNeedForImage = new RawInfoType[]
//            {RawInfoType.TITLE, RawInfoType.CONTENT, RawInfoType.PARA_ABOVE_IMAGE,
//             RawInfoType.PARA_UNDER_IMAGE, RawInfoType.DESC_OF_IMAGE};
    protected RawInfoType[] rawInfoNeedForImage = new RawInfoType[]
        {RawInfoType.CONTENT};

    public TagSuggestModelByER() {
        rawInfoExtractor = new RawInfoExtractorByText();
        featureExtractor = new FeatureExtractorByExpandRank();
        modelKernel = new ModelKernelByWeighted();
    }

    public String trainModel(List<SuggestModelArticleDto> articles) {
        return ((FeatureExtractorByExpandRank)featureExtractor).trainExpandRank(articles);
    }

    public boolean loadModel(String modelPath) {
        try {
            ((FeatureExtractorByExpandRank)featureExtractor).loadExpandRankModel(modelPath);
        } catch (IOException e) {
            e.printStackTrace();
            return false;
        }
        return true;
    }

    public List<SuggestArticleKeywordsDto> suggestTagOfArticle(RecommendArticleDto article) {
        ArrayList<Feature> features = new ArrayList<Feature>();
        for (RawInfoType rawInfoType : rawInfoNeed) {
            RawInfo rawInfo = rawInfoExtractor.extractRawInfo(article,  rawInfoType);
            features.add(featureExtractor.extractFeature(rawInfo, rawInfoType));
        }
        List<TagInfo> tags = modelKernel.suggestTag(features, new RawInfoType[rawInfoNeed.length]);
        List<SuggestArticleKeywordsDto> keywords = new ArrayList<>();
        for (TagInfo tag : tags) {
            keywords.add(new SuggestArticleKeywordsDto(tag.getTagName(), tag.getTagScore()));
        }
        return keywords;
    }

    public List<List<TagInfo>> suggestTagOfImagedArticle(SuggestModelArticleDto article, List<SuggestModelImageDto> imageList) {
        ArrayList<List<TagInfo>> tagsList = new ArrayList<List<TagInfo>>();
        for (SuggestModelImageDto image : imageList) {
            ArrayList<Feature> features = new ArrayList<Feature>();
            for (RawInfoType rawInfoType : rawInfoNeedForImage) {
                RawInfo rawInfo = rawInfoExtractor.extractRawInfo(article, image, rawInfoType);
                features.add(featureExtractor.extractFeature(rawInfo, rawInfoType));
            }
            tagsList.add(modelKernel.suggestTag(features, rawInfoNeedForImage));
        }
        return tagsList;
    }

//    private List<TagInfo> fillTagInfo(List<TagInfo> rawTags) {
//        for (TagInfo tag : rawTags) {
//            tag.setTagID("id-" + UUID.randomUUID().toString());
//            tag.setTagSource("文本抽取");
//        }
//        return rawTags;
//    }
}
